from agent import QueryRewritingAgent
from agent import MultiChunkConclusionAgent
from agent import CitationAgent
from agent import GenerateOutlineAgent
from agent import Agent
import argparse, json
import time
import os
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
import copy
from colorama import Fore, Style


def main():
    parser = argparse.ArgumentParser(description="大模型科研工具V1.0")
    parser.add_argument('--topic', type=str, required=False, help='用户查询', default="explain to me about transformer")
    # parser.add_argument('--model', type=str, required=False, help='基座模型', default="glm-4-airx")
    # parser.add_argument('--excellent_model', type=str, required=False, help='优质模型', default="glm-4-plus")
    parser.add_argument('--model', type=str, required=False, help='基座模型', default="gpt-4o-mini")
    parser.add_argument('--excellent_model', type=str, required=False, help='优质模型', default="gpt-4o-mini")
    parser.add_argument('--token', type=str, required=False, help='API_KEY', default="")
    parser.add_argument('--debug', type=bool, required=False, help='DEBUG模式', default=True)
    parser.add_argument('--output_path', type=str, required=False, help='输出路径', default="review.md")
    parser.add_argument('--conclusion_timeout', type=float, required=False, help='最大chunk总结等待时延', default=30)
    
    args = parser.parse_args()
    if args.token:
        os.environ["ZHIPUAI_API_KEY"] = args.token
    extra_args = {}
    extra_args["excellent_model"] = args.excellent_model
    
    model = args.model
    topic = args.topic
    
    # 开始计时
    start_time = time.time()
    
    # 查询重写，并发查询chunk
    if args.debug:
        print(f">> 用户查询: 【{topic}】")
        print(">> 查询重写:")
    query_rewriting_agent = QueryRewritingAgent(model=model)
    query_list = query_rewriting_agent.rewrite(
        topic=topic,
        num_queries=8,
    )
    if args.debug:
        query_rewriting_agent.PE.print_json(query_list)
        print("-"*100)
    if args.debug:
        print(">> 并发查询chunk:")
    chunk_list = query_rewriting_agent.retrieve(
        query_list=query_list,
        max_return_size=32
    )
    
    if args.debug:
        print(f"chunk_list: (共 {len(chunk_list)} 个结果)\n", json.dumps([
                                                                            f'{ele["id"]}: {ele["entity"]["paper_title"]}' for ele in chunk_list
                                                                        ], indent=2, ensure_ascii=False) + "\n\n...\n")
        print("-"*100)
        
    
    
    # 批量生成总结
    if args.debug:
        print(">> 批量生成总结...")
    multi_agent = MultiChunkConclusionAgent(model=model, DEBUG=args.debug, max_retries=1)
    chunk_list = multi_agent.conclude(
        topic=topic,
        chunk_list=chunk_list,
        return_type="chunk",
        timeout=args.conclusion_timeout,
    )
    
    # 生成大纲
    paper_list = [
        {
            "id": ele["id"],
            "title": ele["entity"]["paper_title"], 
            "总结": ele["总结"]
        } for ele in chunk_list if ("总结" in ele and ele["总结"])
    ]
    if args.debug:
        print(">> 生成大纲:")
        print(">> paper_list（候选 chunk 数量）: ", len(paper_list))
    GO_agent = GenerateOutlineAgent(model=model)
    while True:
        try:
            rough_outline = GO_agent.generate_rough_outline(topic=topic, paper_list=paper_list)
            # rough_outline_title = rough_outline["title"]
            # rough_outline_sections = rough_outline["sections"]
            if args.debug:
                GO_agent.PE.print_json(rough_outline)
                print("-"*100)
            break
        except Exception as e:
            print("[error] 生成大纲失败，重新尝试中...")
            
    
    # 生成二级大纲
    
    def generate_secondary_outline_concurrently(section):
        return generate_secondary_outline(section, topic, rough_outline, model, args, extra_args)
        
    complete_outline = copy.deepcopy(rough_outline)
    # 使用 ThreadPoolExecutor 进行并发处理
    with ThreadPoolExecutor(max_workers=128) as executor:
        # 提交所有任务
        future_to_section = {executor.submit(generate_secondary_outline_concurrently, section): section for section in rough_outline["sections"]}

        for future in as_completed(future_to_section):
            section, secondary_outline = future.result()
            for i in range(len(complete_outline["sections"])):
                if complete_outline["sections"][i]["name"] == section["name"]:
                    complete_outline["sections"][i]["subsections"] = secondary_outline["subsections"]
                    for j in range(len(complete_outline["sections"][i]["subsections"])):
                        complete_outline["sections"][i]["subsections"][j]["context info"] = {}
                        if j > 0:
                            complete_outline["sections"][i]["subsections"][j]["context info"]["current section info"] = {
                                "section name": complete_outline["sections"][i]["name"],
                                "section description": complete_outline["sections"][i]["description"],
                            }
                            complete_outline["sections"][i]["subsections"][j]["context info"]["former subsection name"] = complete_outline["sections"][i]["subsections"][j - 1]["name"]
                        if j < len(complete_outline["sections"][i]["subsections"]) - 1:
                            complete_outline["sections"][i]["subsections"][j]["context info"]["next subection name"] = complete_outline["sections"][i]["subsections"][j + 1]["name"]
    
    GO_agent.PE.print_json(complete_outline)
    
    
    final_outline = complete_outline
    
    # 为每个章节生成综述内容
    total_raw_citation = {}
    total_chunk_list = []

    def process_section_concurrently(section, subsection):
        return process_section(section, subsection, topic, final_outline, model, args, extra_args)

    # 使用 ThreadPoolExecutor 进行并发处理
    with ThreadPoolExecutor(max_workers=128) as executor:
        
        future_to_section = {
            executor.submit(process_section_concurrently, section, subsection): subsection
            for section in final_outline["sections"]
            for subsection in section["subsections"]
        }

        for future in as_completed(future_to_section):
            section, subsection, raw_citation, current_chunk_list = future.result()
            if section["name"] not in total_raw_citation:
                total_raw_citation[section["name"]] = {}
            print("-"*100+"\n"+"-"*100+"\n"+"-"*100)
            print('>> total_raw_citation:')
            GO_agent.PE.print_json(total_raw_citation)
            print('>> section:')
            GO_agent.PE.print_json(section)
            print('>> subsection:')
            GO_agent.PE.print_json(subsection)
            total_raw_citation[section["name"]][subsection["name"]] = raw_citation
            total_chunk_list.extend(current_chunk_list)
    
    citation_agent = CitationAgent(model=model)
    
    final_total_raw_citation = f'# {final_outline["title"]}\n\n'
    
    for section in final_outline["sections"]:
        final_total_raw_citation += f'## {section["name"]}\n\n'
        for subsection in section["subsections"]:
            final_total_raw_citation += f'### {subsection["name"]}\n\n{total_raw_citation[section["name"]][subsection["name"]]}\n\n'
    
    formatted_citation, reference_suffix = citation_agent.format_citation(raw_citation=final_total_raw_citation, chunk_list=total_chunk_list)
    final_citation = formatted_citation + "\n\n" + reference_suffix
    print(">> 综述生成结果：")
    print("-"*100+"\n"+"-"*100+"\n"+"-"*100+"\n")
    print(final_citation)
    
    # 将 final_citation 输出到文件
    with open(args.output_path, 'w', encoding='utf-8') as file:
        file.write(final_citation)
    
    # 结束计时
    end_time = time.time()
    execution_time = end_time - start_time
    print(f"程序执行时间: {Fore.GREEN}{execution_time:.2f}{Style.RESET_ALL} 秒")
    
def generate_secondary_outline(section, topic, rough_outline, model, args, extra_args):
    if args.debug:
        print(">> [generate_secondary_outline] ...")
        print("-"*100+"\n"+"-"*100+"\n"+"-"*100+"\n")
    global_agent = Agent()
    section_topic = topic+" "+section["name"]+" "+section["description"]
    query_rewriting_agent = QueryRewritingAgent(model=model)
    query_list = query_rewriting_agent.rewrite(
        topic=section_topic,
        num_queries=8,
    )

    chunk_list = query_rewriting_agent.retrieve(
        query_list=query_list,
        max_return_size=16
    )
    multi_agent = MultiChunkConclusionAgent(model=model, DEBUG=True, max_retries=1, timeout=10)
    chunk_list = multi_agent.conclude(
        topic=section_topic,
        chunk_list=chunk_list,
        return_type="chunk",
        return_relevant=True,
        timeout=args.conclusion_timeout * 0.85,
    )
    
    simplified_chunk_list = [
        {"id": ele["id"], "总结": ele["总结"]} for ele in chunk_list if "总结" in ele
    ]
    GO_agent = GenerateOutlineAgent(model=model)
    secondary_outline = GO_agent.generate_secondary_outline(topic, section, simplified_chunk_list, rough_outline)
    return section, secondary_outline
    

def process_section(section, subsection, topic, rough_outline, model, args, extra_args):
    if args.debug:
        print(">> [process_section] ...")
        print("-"*100+"\n"+"-"*100+"\n"+"-"*100+"\n")
        
    global_agent = Agent()
    section_simple = {"name": section["name"], "description": section["description"]}
    # current_topic = f'原始意图: {topic}\n综述标题（拟）: {rough_outline["title"]}\n当前章节: {global_agent.PE.as_json(section_simple)}\n当前小节：{global_agent.PE.as_json(subsection)}\n'
    current_topic = f'{global_agent.PE.as_json(subsection)}\n'
    query_rewriting_agent = QueryRewritingAgent(model=model)
    query_list = query_rewriting_agent.rewrite(
        topic=current_topic,
        num_queries=8,
    )
    chunk_list = query_rewriting_agent.retrieve(
        query_list=query_list,
        max_return_size=16
    )
    if args.debug:
        print(">> [section] 批量生成总结...")
    
    multi_agent = MultiChunkConclusionAgent(model=model, DEBUG=args.debug, max_retries=1)
    chunk_list = multi_agent.conclude(
        topic=current_topic,
        chunk_list=chunk_list,
        return_type="chunk",
        return_relevant=True,
        timeout=args.conclusion_timeout,
    )
    if args.debug:
        print(f"chunk_list: (共 {len(chunk_list)} 个结果)\n", json.dumps(chunk_list, indent=2, ensure_ascii=False)[:256] + "\n\n...\n")
        print("-"*100)
        
    cite_agent = CitationAgent(model=extra_args["excellent_model"], DEBUG=args.debug)
    raw_citation = cite_agent.cite_section(
        section_topic=current_topic,
        chunk_list=chunk_list,
        DEBUG=args.debug,
        enable_citation_formatting=False,
    )
    if "Request timed out" in raw_citation:
        cite_agent = CitationAgent(model=model, DEBUG=args.debug)
        raw_citation = cite_agent.cite_section(
            section_topic=current_topic,
            chunk_list=chunk_list,
            DEBUG=args.debug,
            enable_citation_formatting=False,
        )
    raw_citation = raw_citation.replace("```", "")
    if args.debug:
        print(">> raw_citation:")
        print("-"*100+"\n"+"-"*100+"\n"+"-"*100+"\n")
        print(raw_citation)
        print("-"*100+"\n"+"-"*100+"\n"+"-"*100+"\n")
    return section, subsection, raw_citation, chunk_list
        
    

if __name__=="__main__":
    main()